import os
import json
from typing import List, Dict, Optional
from langchain_deepseek import chat_models
from outline_gen_agents.uilts.outlines_state import OutLinesState


def _get_provider_conf_from_env(provider: str) -> Optional[Dict[str, str]]:
    key = provider.upper()
    api_key = os.getenv(f"{key}_API_KEY")
    base_url = os.getenv(f"{key}_BASE_URL")
    model = os.getenv(f"{key}_MODEL")
    if api_key and base_url and model:
        return {"api_key": api_key, "base_url": base_url, "model": model}
    return None


def agent1_requirements(state: OutLinesState) -> OutLinesState:
    user_input = state.get("user_input")
    provider = state.get("provider")
    system_prompt = "你是一个科学计量报告的专业撰写人员，可以非常准确地理解用户的需求，并能将其中的具体需求进行格式化抽取和识别。"
    user_prompt = f"""
    请阅读用户进行的输入，抽取并返回一个JSON对象，包含字段：
    - user_need: 用户的核心需求，简洁描述，返回一个字符串列表，每个元素为一个具体的需求。
    - goal: 分析与报告的目标，返回一个字符串。
    - audience: 面向的受众，返回一个字符串，如'政府决策者'、'研究人员'、'企业'等。
    - domain: 具体技术/行业领域，返回一个字符串，如'人工智能'、'量子科技'等。
    - year: 关注年份，使用整数类型返回具体的年份列表，如[2023, 2024]、[2023]等。
    - data: 与需求相关的数据源（列表），如['专利'、'论文'、'技术报告']
    - analysis_paradigm：用户需求呈现的分析范式，返回一个字符串，字符串只能是列表['宏观态势','前沿监测','绩效评估','其他']中的一个
    用户需求：
    {user_input}
    注意：
    - 仅返回可解析的JSON字符串，不要包含解释或额外文本。
    """

    conf = _get_provider_conf_from_env(provider)
    llm = chat_models.ChatDeepSeek(model=conf["model"], base_url=conf["base_url"], api_key=conf["api_key"],
                                   model_kwargs={"response_format": {"type": "json_object"}}
                                   )
    msg = llm.invoke(
        [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    parsed = json.loads(msg.content.strip())
    state["messages"].append(msg)
    state.update(parsed)
    return state


def node_macro(state: OutLinesState) -> OutLinesState:
    system_prompt = f"""
    你是一个使用科学计量方法撰写宏观态势报告大纲的专业撰写人员，你需要根据用户输入的进行报告大纲的撰写。
    用户输入信息包含如下字段，请着重进行理解：
    - user_need: 用户的核心需求，这是一个字符串列表，每个元素为一个具体的需求
    - goal: 分析与报告的目标
    - audience: 面向的受众
    - domain: 具体技术/行业领域
    - year: 关注年份
    - data: 数据源
    你需要以json形式返回一个大纲列表，大纲中需要尽可能详细的列出所有的需求，包含字段：
    - outlines: 用户要撰写的大纲，返回一个字符串列表，每个元素为一个大纲的章节。
    注意：
    - 仅返回可解析的JSON字符串，不要包含解释或额外文本。
    - 其中，每一个大纲的章节请保证具有可操作性，限定性，具备分析方法。   
    """
    user_prompt = f"""
    - user_need: {state["user_need"]}
    - goal: {state["goal"]}
    - audience: {state["audience"]}
    - domain: {state["domain"]}
    - year: {state['year']}
    - data: {state['data']}
    """

    conf = _get_provider_conf_from_env(state['provider'])
    llm = chat_models.ChatDeepSeek(model=conf["model"], base_url=conf["base_url"], api_key=conf["api_key"],
                                   model_kwargs={"response_format": {"type": "json_object"}}
                                   )
    msg = llm.invoke(
        [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    parsed = json.loads(msg.content.strip())
    state["messages"].append(msg)
    state.update(parsed)
    return state


def node_frontier(state: OutLinesState) -> OutLinesState:
    # 前沿监测
    system_prompt = f"""
        你是一个使用科学计量方法撰写前沿监测报告大纲的专业撰写人员，你需要根据用户输入的进行报告大纲的撰写。
        用户输入信息包含如下字段，请着重进行理解：
        - user_need: 用户的核心需求，这是一个字符串列表，每个元素为一个具体的需求
        - goal: 分析与报告的目标
        - audience: 面向的受众
        - domain: 具体技术/行业领域
        - year: 关注年份
        - data: 数据源
        你需要以json形式返回一个大纲列表，大纲中需要尽可能详细的列出所有的需求，包含字段：
        - outlines: 用户要撰写的大纲，返回一个字符串列表，每个元素为一个大纲的章节。
        注意：
        - 仅返回可解析的JSON字符串，不要包含解释或额外文本。
        - 其中，每一个大纲的章节请保证具有可操作性，限定性，具备分析方法。   
        """
    user_prompt = f"""
        - user_need: {state["user_need"]}
        - goal: {state["goal"]}
        - audience: {state["audience"]}
        - domain: {state["domain"]}
        - year: {state['year']}
        - data: {state['data']}
        """

    conf = _get_provider_conf_from_env(state['provider'])
    llm = chat_models.ChatDeepSeek(model=conf["model"], base_url=conf["base_url"], api_key=conf["api_key"],
                                   model_kwargs={"response_format": {"type": "json_object"}})

    msg = llm.invoke(
        [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    parsed = json.loads(msg.content.strip())
    state["messages"].append(msg)
    state.update(parsed)
    return state


def node_performance(state: OutLinesState) -> OutLinesState:
    system_prompt = f"""
        你是一个使用科学计量方法撰写绩效评估报告大纲的专业撰写人员，你需要根据用户输入的进行报告大纲的撰写。
        用户输入信息包含如下字段，请着重进行理解：
        - user_need: 用户的核心需求，这是一个字符串列表，每个元素为一个具体的需求
        - goal: 分析与报告的目标
        - audience: 面向的受众
        - domain: 具体技术/行业领域
        - year: 关注年份
        - data: 数据源
        你需要以json形式返回一个大纲列表，大纲中需要尽可能详细的列出所有的需求，包含字段：
        - outlines: 用户要撰写的大纲，返回一个字符串列表，每个元素为一个大纲的章节。
        注意：
        - 仅返回可解析的JSON字符串，不要包含解释或额外文本。
        - 其中，每一个大纲的章节请保证具有可操作性，限定性，具备分析方法。   
        """
    user_prompt = f"""
        - user_need: {state["user_need"]}
        - goal: {state["goal"]}
        - audience: {state["audience"]}
        - domain: {state["domain"]}
        - year: {state['year']}
        - data: {state['data']}
        """

    conf = _get_provider_conf_from_env(state['provider'])
    llm = chat_models.ChatDeepSeek(model=conf["model"], base_url=conf["base_url"], api_key=conf["api_key"],
                                   model_kwargs={"response_format": {"type": "json_object"}}
)
    msg = llm.invoke(
        [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    parsed = json.loads(msg.content.strip())
    state["messages"].append(msg)
    state.update(parsed)
    return state


def node_other(state: OutLinesState) -> OutLinesState:
    # 其他
    state["analysis_paradigm"] = "其他"
    return state


def route_paradigm(state: OutLinesState) -> str:
    # 根据 analysis_paradigm 路由到对应节点
    val = (state.get("analysis_paradigm") or "").strip()
    mapping = {
        "宏观态势报告": "宏观态势",
        "宏观态势": "宏观态势",
        "前沿监测报告": "前沿监测",
        "前沿监测": "前沿监测",
        "绩效评估报告": "绩效评估",
        "绩效评估": "绩效评估",
        "其他": "其他",
    }
    return mapping.get(val, "其他")
